package org.apache.flink.streaming.api.datastream;

import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.metrics.Counter;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

import java.util.TreeSet;

@SuppressWarnings("unused")
public abstract class KeyedWatermarkProcessFunction<K, I, O> extends KeyedProcessFunction<K, I, O> {
    ValueState<Long> watermark;
    @SuppressWarnings("rawtypes")
    ValueState<TreeSet> event;
    Counter lateElement, globalLateElement;

    @Override
    public void open(Configuration parameters) {
        watermark = getRuntimeContext().getState(new ValueStateDescriptor<>("Watermark", Long.class));
        event = getRuntimeContext().getState(new ValueStateDescriptor<>("Event", TreeSet.class));
        lateElement = getRuntimeContext().getMetricGroup().counter("LateElement");
        globalLateElement = getRuntimeContext().getMetricGroup().counter("GlobalLateElement");
    }

    @Override
    public void processElement(I value, Context ctx, Collector<O> out) throws Exception {
        boolean late = false;
        if (watermark.value() != null) {
            late = detectLate(value);
            detectGlobalLate(value, ctx);
        }
        if (!late) {
            update(value);
            triggerEvent(ctx.getCurrentKey());
            process(value, ctx, out);
        }
    }

    public abstract void process(I value, Context ctx, Collector<O> out) throws Exception;

    public abstract long extract(I value);

    private void update(I value) throws Exception {
        watermark.update(extract(value));
    }

    private boolean detectLate(I value) throws Exception {
        if (extract(value) < watermark.value()) {
            lateElement.inc();
            processLate(value);
            return true;
        }
        return false;
    }

    public void processLate(I value) {
    }

    private void detectGlobalLate(I value, Context ctx) {
        if (extract(value) < ctx.timerService().currentWatermark()) {
            globalLateElement.inc();
            processGlobalLate(value);
        }
    }

    public void processGlobalLate(I value) {
    }

    private void triggerEvent(K key) throws Exception {
        if (event.value() != null) {
            @SuppressWarnings("unchecked")
            TreeSet<Long> set = (TreeSet<Long>) event.value();
            while (!set.isEmpty() && set.first() <= watermark.value()) {
                // noinspection DataFlowIssue
                execute(set.pollFirst(), key);
            }
        }
    }

    public void register(long timestamp) throws Exception {
        if (event.value() == null) {
            event.update(new TreeSet<Long>());
        }
        @SuppressWarnings("unchecked")
        TreeSet<Long> set = (TreeSet<Long>) event.value();
        set.add(timestamp);
    }

    public abstract void execute(long timestamp, K key) throws Exception ;
}
